/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2022. All rights reserved.
 * Description: hdds op remote look up execute.
 * Author: qingjicheng
 * Create: 2022-11-28
 */

#include "op_hdcs_remote_lookup_executor.h"
#include "transport_shm_event_pub.h"
#include "log.h"
#include "adapter_hdds.h"
#include "adapter_hal.h"
#include "hccl/hccl_ex.h"
#include "../../../../hccl_heterog/rpc/hccl_comm_rpc.h"
#include "../hccl_heterog/rpc/hcom_rpc.h"

using namespace std;

namespace hccl {
constexpr u32 PS_INVALID_VALUE = 65536;
constexpr u64 PRESET_KEY_SIZE = 20 * 1024 * 1024;
// 性能优化调参
constexpr u32 EACH_PS_KEY_NUM = 10;

HdcsOpRemoteLookUpExecutor::HdcsOpRemoteLookUpExecutor() : taskExeFuncTable_({
    {OpTaskType::SEND, std::bind(&HdcsOpRemoteLookUpExecutor::TaskSendKeys, this, std::placeholders::_1)},
    {OpTaskType::RECV, std::bind(&HdcsOpRemoteLookUpExecutor::TaskWaitLookupFinish, this, std::placeholders::_1)},
    {OpTaskType::DROP_DUPLICATES, std::bind(&HdcsOpRemoteLookUpExecutor::TaskDropDuplicates, this,
        std::placeholders::_1)},
    {OpTaskType::KEYS_DUPLICATES, std::bind(&HdcsOpRemoteLookUpExecutor::KeyReduce, this,
        std::placeholders::_1)},
    {OpTaskType::RECOVER_VALUE, std::bind(&HdcsOpRemoteLookUpExecutor::ParralAicpuRecoverValue, this,
        std::placeholders::_1)},
    {OpTaskType::WAIT_SEND_KEY_FINISH, std::bind(&HdcsOpRemoteLookUpExecutor::WaitSendKeyFinish, this,
        std::placeholders::_1)},
    {OpTaskType::RESET_UNIQUE_TASK, std::bind(&HdcsOpRemoteLookUpExecutor::TaskRestUnique, this,
        std::placeholders::_1)},
})
{}

HdcsOpRemoteLookUpExecutor::~HdcsOpRemoteLookUpExecutor()
{
    CHK_PRT(DeInit());
}

HcclResult HdcsOpRemoteLookUpExecutor::Init(std::vector<u32> &psRankId, u64 pipelineKeyNum)
{
    HCCL_INFO("[HdcsOpRemoteLookUpExecutor][Init] pipelineKeyNum[%llu], psRankId[%u].", pipelineKeyNum, psRankId[0]);
    psRankId_ = psRankId;

    pipelineKeyNum_ = pipelineKeyNum;
    psSize_ = psRankId_.size();

    // psId应该来源于roleTable，而不是在这里写死，这里有可能rankId与psId的映射关系是错的
    for (size_t index = 0; index < psSize_; index++) {
        rankId2PsId_[psRankId_[index]] = index;
    }

#ifdef CCL_KERNEL
    shardExecMaxCoreNum_ = GetCPUNum();
    HCCL_INFO("shardExecMaxCoreNum_[%u].", shardExecMaxCoreNum_);
    keyPsSizeAfterUnique_.resize(shardExecMaxCoreNum_);
    psKeyNumForShard_.resize(shardExecMaxCoreNum_);
    psOffsetForShard_.resize(shardExecMaxCoreNum_);
    threadkeyNumForShard_.resize(shardExecMaxCoreNum_);
    recoverValueParas_.resize(shardExecMaxCoreNum_);
    prepareSendKeysBufParas_.resize(shardExecMaxCoreNum_);

    for (u32 i = 0; i < shardExecMaxCoreNum_; ++i) {
        psKeyNumForShard_[i].resize(psSize_, 0);
        psOffsetForShard_[i].resize(psSize_, 0);
        keyPsSizeAfterUnique_[i].resize(psSize_);
        threadkeyNumForShard_[i].reserve(PRESET_KEY_SIZE);
    }

#endif

    return HCCL_SUCCESS;
}

HcclResult HdcsOpRemoteLookUpExecutor::DeInit()
{
    if (psIdsMem_ != nullptr) {
        delete[] psIdsMem_;
        psIdsMem_ = nullptr;
    }
    return HCCL_SUCCESS;
}

HcclResult HdcsOpRemoteLookUpExecutor::ExecuteOp(HcclOpPtr opPtr)
{
    CHK_PTR_NULL(opPtr);
    HcclOperator<EmbeddingServiceParam> *ptr = static_cast<HcclOperator<EmbeddingServiceParam> *>(opPtr);

    // 给partMap查询结果申请内存
    if ((ptr->param_.embeddingParam.cubeIndex <= 0) &&
        (ptr->param_.keyMaxNum > keyMaxNum_)) {
        if (psIdsMem_ != nullptr) {
            delete[] psIdsMem_;
            psIdsMem_ = nullptr;
        }

        keyMaxNum_ = ptr->param_.keyMaxNum;
        psIdsMem_ = new (std::nothrow) u32[keyMaxNum_];
        CHK_PTR_NULL(psIdsMem_);
        HCCL_INFO("lookup tag[%d] alloc new mem, size[%u]", ptr->param_.tag, keyMaxNum_);
    }

    CHK_RET(ExecuteTask(ptr));

    return HCCL_SUCCESS;
}

HcclResult HdcsOpRemoteLookUpExecutor::ExecuteTask(HcclOperator<EmbeddingServiceParam> *opPtr)
{
    auto &taskQue = GetTaskQue<EmbeddingServiceParam>(opPtr);
    while (!taskQue.empty()) {
        HCCL_INFO("ExecuteTask[%d].", taskQue.front()->taskType);
        CHK_RET(taskExeFuncTable_[taskQue.front()->taskType](taskQue.front()));
        taskQue.pop();
    }
    return HCCL_SUCCESS;
}

HcclResult HdcsOpRemoteLookUpExecutor::TaskDropDuplicates(OpTaskPtr &task)
{
    CHK_RET(TaskKeysDuplicatesShard(task));
    CHK_RET(TaskSendKeys(task));
    CHK_RET(TaskWaitLookupFinish(task));
    CHK_RET(TaskRecoverValueShard(task));

    return HCCL_SUCCESS;
}

HcclResult HdcsOpRemoteLookUpExecutor::TaskKeysDuplicatesShard(OpTaskPtr &task)
{
    return HCCL_SUCCESS;
}

HcclResult HdcsOpRemoteLookUpExecutor::ReduceKeyAndGetMappingMatrix(PartitionMapTask &params)
{
    u32 keyNum = params.keyMaxNum;
    if (params.uniqued) {
        CHK_PRT_RET((*params.keyNumInput > UINT32_MAX),
            HCCL_ERROR("keyNumInput[%lld] > UINT32_MAX", *params.keyNumInput), HCCL_E_PARA);

        keyNum = static_cast<u32>(*params.keyNumInput);
    }

    ConUniqueMappersInfo<s64, s32> info{};
    info.conMapperNum = psSize_;
    info.parallelNum = CalThreadNumEachPs(keyNum);
    info.maxParallelNum = shardExecMaxCoreNum_;
    info.srcData.eles = params.keys;
    info.srcData.count = keyNum;
    info.srcData.elesCapacity = params.keyMaxNum;
    info.uniqueData.eles = static_cast<s64 *>(params.keyTransferMem);
    info.uniqueData.count = 0;
    info.uniqueData.elesCapacity = params.keyMaxNum;
    info.preSlicingOffsetAndNum = psOffsetAndNum_.Get();
    info.preSlicingInfo = psIdsMem_;
    info.mapReserveNum = keyNum / psSize_ / info.parallelNum;
    info.splitData = true;
    info.enableKeyCounter = params.enableKeyCounter;

    TIME_PRINT(CHK_RET(conUniqueMappersMgr_.Cfg(info)));

    Buffer<s32> mappingMatrixS32;
    mappingMatrixS32.eles = mappingMatrix_.Get();
    mappingMatrixS32.count = 0;
    mappingMatrixS32.elesCapacity = params.keyMaxNum;

    TIME_PRINT(CHK_RET(conUniqueMappersMgr_.UniqueDataAndGetMappingMatrix(mappingMatrixS32)));

    CHK_RET(FetchUniquedCounts());

    PRINT_ARRAY(static_cast<float *>(params.value),
        params.keyMaxNum * params.valueItemSize / sizeof(float), "LookupValues");

    return HCCL_SUCCESS;
}

HcclResult HdcsOpRemoteLookUpExecutor::FillCounterWhileUniqued(PartitionMapTask &partMapTask)
{
    CHK_PTR_NULL(partMapTask.keyCount);

    s64 numUniqued = *(static_cast<s64 *>(partMapTask.numUniqued));

    // 当前ge是按照s64传入的
    s64 *keyCount = reinterpret_cast<s64 *>(partMapTask.keyCount);
    KeyAndCounter<s64, s64> *keyTransMem = static_cast<KeyAndCounter<s64, s64> *>(partMapTask.keyTransferMem);

    if (psSize_ == OpEmbeddingServiceExecutor::PS_SIZE_ONE) {
        for (s64 i = 0; i < numUniqued; i++) {
            keyTransMem[i].key = partMapTask.keys[i];
            keyTransMem[i].counter = keyCount[i];
        }
    } else {
        for (s64 i = 0; i < numUniqued; i++) {
            // 前序流程已经将key放入发送内存，且映射矩阵已经ready
            s32 &index = mappingMatrix_[i];
            keyTransMem[index].counter = keyCount[i];
        }
    }

    PRINT_ARRAY(keyCount, numUniqued, "LookupKeyCount");

    return HCCL_SUCCESS;
}

HcclResult HdcsOpRemoteLookUpExecutor::KeyReduceWithUniqued(PartitionMapTask &partMapTask)
{
    s64 numUniqued = *(static_cast<s64 *>(partMapTask.numUniqued));

    // 1个PS的情况
    if (psSize_ == OpEmbeddingServiceExecutor::PS_SIZE_ONE) {
        // 特征准入场景
        if (partMapTask.enableKeyCounter) {
            HCCL_DEBUG("Host Uniqued Single PS: KeyCounter is enabled");
            CHK_RET(FillCounterWhileUniqued(partMapTask));
        } else {
            s64 len = numUniqued * sizeof(s64);
            CHK_SAFETY_FUNC_RET(memcpy_s(partMapTask.keyTransferMem, len, partMapTask.keys, len));
        }

        // indices当前为s32
        s64 len = partMapTask.keyMaxNum * sizeof(s32);
        mappingMatrix_.Set(static_cast<s32 *>(partMapTask.indices));
        CHK_SAFETY_FUNC_RET(memcpy_s(mappingMatrix_.Get(), len, partMapTask.uniqueIndices, len));

        psOffsetAndNum_[0].offset = 0;
        psOffsetAndNum_[0].count = numUniqued;
        HCCL_DEBUG("Lookup KeyReduce uniqued 1 PS executed");
    } else {
        // uniqued也走ReduceKey实现多线程搬移数据
        TIME_PRINT(CHK_RET(ReduceKeyAndGetMappingMatrix(partMapTask)));
        PRINT_ARRAY(mappingMatrix_.Get(), numUniqued, "Lookup Uniqued Partmap MappingMatrix");

        HcclResult ret = ConUniqueMappersMgr<s64, s32>::ParallelConvertMappingMatrixEqually(mappingMatrix_.Get(),
            numUniqued, partMapTask.uniqueIndices, static_cast<u32>(partMapTask.keyMaxNum),
            static_cast<s32 *>(partMapTask.indices), shardExecMaxCoreNum_);
        CHK_PRT_RET(ret != HCCL_SUCCESS,
            HCCL_ERROR("ParallelConvertMappingMatrixEqually failed, ret[%d]", ret), ret);

        if (partMapTask.enableKeyCounter) {
            HCCL_DEBUG("Host Uniqued Multi PS: KeyCounter is enabled");
            CHK_RET(FillCounterWhileUniqued(partMapTask));
        }

        // 将mappingMatrix_指向最终生成的恢复矩阵
        mappingMatrix_.Set(static_cast<s32 *>(partMapTask.indices));
        HCCL_DEBUG("Lookup KeyReduce uniqued multi PS executed");
    }

    return HCCL_SUCCESS;
}

HcclResult HdcsOpRemoteLookUpExecutor::KeyReduce(OpTaskPtr &task)
{
    PartitionMapTask *partMapTask = dynamic_cast<PartitionMapTask *>(task.get());

    if (partMapTask->usePipeline && partMapTask->cubeIndex > 0) {
        // 需要做流水的reduceKey，并更新psOffsetAndNum_，待补充
        return HCCL_SUCCESS;
    }

    if (partMapTask->disableUnique && (partMapTask->intZerocpyFlag == ZERO_COPY_USED)) {
        partMapTask->keys = reinterpret_cast<s64 *>(*(partMapTask->keys));
    }

    HCCL_DEBUG("CollRemoteLookUp KeyReduce disableUnique[%d], intZerocpyFlag[%d]",
        partMapTask->disableUnique, partMapTask->intZerocpyFlag);

    PrepareMappingInfo(*partMapTask);
    TIME_PRINT(CHK_RET(GetPsIdInfo(*partMapTask)));

    if (partMapTask->disableUnique) {
        // 考虑对单PS的流程进一步优化
        TIME_PRINT(CHK_RET(ReduceKeyAndGetMappingMatrixWithNoUnique(*partMapTask)));
    } else {
        if (partMapTask->uniqued) {
            TIME_PRINT(CHK_RET(KeyReduceWithUniqued(*partMapTask)));
        } else {
            TIME_PRINT(CHK_RET(ReduceKeyAndGetMappingMatrix(*partMapTask)));
        }
    }

#ifdef DUMP_DATA
    u32 count = partMapTask->keyMaxNum;
    if (partMapTask->enableKeyCounter) {
        count = count * DOUBLE;
    }

    PRINT_ARRAY(reinterpret_cast<s64 *>(psOffsetAndNum_.Get()), psSize_ * DOUBLE, "LookupPsOffset");
    PRINT_ARRAY(mappingMatrix_.Get(), partMapTask->keyMaxNum, "LookupMappingMatrix");
    PRINT_ARRAY(static_cast<s64 *>(partMapTask->keyTransferMem), count, "LookupkeyTransferMem");
#endif

    return HCCL_SUCCESS;
}

void HdcsOpRemoteLookUpExecutor::PrepareMappingInfo(PartitionMapTask &partMapTask)
{
    if (!partMapTask.pairedMode) {
        void *addr = static_cast<s64 *>(partMapTask.keyTransferMem) + partMapTask.keyMaxNum;
        if (partMapTask.enableKeyCounter) {
            addr = static_cast<s8 *>(addr) + static_cast<u64>(partMapTask.keyMaxNum) * ES_KEY_COUNTER_MEM_BYTES_SIZE;
        }

        mappingMatrix_.Set(static_cast<s32 *>(addr));
        psOffsetAndNum_.Resize(psSize_);
        return;
    }

    s64 &numUniqued = *(static_cast<s64 *>(partMapTask.numUniqued));

    // unqiued模式，使用中间workerspace内存，最后输出到外部indices内存！！！
    if (partMapTask.uniqued) {
        void *addr = static_cast<s64 *>(partMapTask.keyTransferMem) + partMapTask.keyMaxNum;
        if (partMapTask.enableKeyCounter) {
            addr = static_cast<s8 *>(addr) + static_cast<u64>(partMapTask.keyMaxNum) * ES_KEY_COUNTER_MEM_BYTES_SIZE;
        }

        mappingMatrix_.Set(static_cast<s32 *>(addr));
        numUniqued = *partMapTask.keyNumInput;
        PRINT_ARRAY(partMapTask.uniqueIndices, partMapTask.keyMaxNum, "UniqueIndices");
    } else {
        mappingMatrix_.Set(static_cast<s32 *>(partMapTask.indices));
        // keyMaxNum为u32，强转无截断问题
        numUniqued = static_cast<s64>(partMapTask.keyMaxNum);
    }

    HCCL_DEBUG("uniqued[%d], numUniqued[%lld]", partMapTask.uniqued, numUniqued);
    PRINT_ARRAY(reinterpret_cast<s64 *>(partMapTask.keys), numUniqued, "Lookup keys");

    // 类型转换确认不存在截断问题
    *(static_cast<s64 *>(partMapTask.psSegNum)) = static_cast<s64>(psSize_);
    psOffsetAndNum_.Set(static_cast<PsBufferInfo *>(partMapTask.psSeg));
}

HcclResult HdcsOpRemoteLookUpExecutor::ParralGetGetBatchPsIds(PartitionMapTask &partMapTask)
{
#ifdef CCL_KERNEL
    // keyMaxNum为u32，强转无截断问题
    s64 keyNum = partMapTask.uniqued ? (*partMapTask.keyNumInput) : static_cast<s64>(partMapTask.keyMaxNum);

    int64_t totalSize = keyNum;
    int64_t unitSize = keyNum / shardExecMaxCoreNum_;
    if ((keyNum % shardExecMaxCoreNum_) != 0) {
        unitSize++;
    }

    ParallelFor(totalSize, unitSize, [&partMapTask, this] (int64_t start, int64_t end) -> HcclResult {
        u64 *srcKeys = reinterpret_cast<u64 *>(partMapTask.keys + start);
        u32 *dst = psIdsMem_ + start;

        int64_t size = end - start;
        CHK_PRT_RET(UNLIKELY(size > UINT32_MAX),
            HCCL_ERROR("size[%lld] > UINT32_MAX, does not support", size), HCCL_E_PARA);
        CHK_PRT_RET(UNLIKELY(size == 0), HCCL_WARNING("size is 0 before hrtGetBatchPsIds"), HCCL_SUCCESS);
        // 查表
        CHK_RET(hrtGetBatchPsIds(srcKeys, &dst, static_cast<u32>(size)));

        return HCCL_SUCCESS;
    });
#endif

    return HCCL_SUCCESS;
}

HcclResult HdcsOpRemoteLookUpExecutor::GetPsIdInfo(PartitionMapTask &partMapTask)
{
    if (psSize_ == OpEmbeddingServiceExecutor::PS_SIZE_ONE) {
        psOffsetAndNum_[0].offset = 0;
        return HCCL_SUCCESS;
    }

    CHK_RET(ParralGetGetBatchPsIds(partMapTask));

    // keyMaxNum为u32，强转无截断问题
    s64 keyNum = partMapTask.uniqued ? (*partMapTask.keyNumInput) : static_cast<s64>(partMapTask.keyMaxNum);

    // 计算主线程sendbuf分块
    u64 offset = 0;
    for (u32 i = 0; i < psSize_ - 1; ++i) {
        psOffsetAndNum_.Get()[i].offset = offset;
        for (s64 keyIndex = 0; keyIndex < keyNum; ++keyIndex) {
            if (psIdsMem_[keyIndex] == i) {
                offset++;
            }
        }
    }

    psOffsetAndNum_.Get()[psSize_ - 1].offset = offset;

    return HCCL_SUCCESS;
}

HcclResult HdcsOpRemoteLookUpExecutor::ReduceKeyAndGetMappingMatrixWithNoUnique(PartitionMapTask &partMapTask)
{
    if (psSize_ == OpEmbeddingServiceExecutor::PS_SIZE_ONE) {
        s64 len = partMapTask.keyMaxNum * sizeof(s64);
        CHK_SAFETY_FUNC_RET(memcpy_s(partMapTask.keyTransferMem, len, partMapTask.keys, len));
#ifdef CCL_KERNEL

        int64_t totalSize = partMapTask.keyMaxNum;
        int64_t unitSize = totalSize / shardExecMaxCoreNum_;
        if ((totalSize % shardExecMaxCoreNum_) != 0) {
            unitSize++;
        }

        ParallelFor(totalSize, unitSize, [this] (int64_t start, int64_t end) -> HcclResult {
            for (int64_t i = start; i < end; i++) {
                mappingMatrix_[i] = static_cast<s32>(i);
            }

            return HCCL_SUCCESS;
        });
#endif

        psOffsetAndNum_[0].count = partMapTask.keyMaxNum;

        return HCCL_SUCCESS;
    }

    u32 counts[psSize_] = { 0 };
    for (u32 i = 0; i < partMapTask.keyMaxNum; i++) {
        u32 psId = psIdsMem_[i];

        s64 index = psOffsetAndNum_[psId].offset + counts[psId];
        CHK_PRT_RET(UNLIKELY(index > INT32_MAX),
            HCCL_ERROR("index[%lld] > INT32_MAX, current does not support", index), HCCL_E_NOT_SUPPORT);

        // 受aicore算子限制，恢复矩阵当前只支持s32
        mappingMatrix_[i] = static_cast<s32>(index);
        s64 *dstMem = static_cast<s64 *>(partMapTask.keyTransferMem) + index;
        *dstMem = partMapTask.keys[i];
        counts[psId]++;
    }

    for (u32 i = 0; i < psSize_; i++) {
        psOffsetAndNum_[i].count = counts[i];
    }

    return HCCL_SUCCESS;
}

u32 HdcsOpRemoteLookUpExecutor::CalThreadNumEachPs(u32 keyMaxNum)
{
    // 至少一个线程
    u32 threadNumEachPs = 1;
    u32 eachPsKeyNum = keyMaxNum / psSize_;

    if (eachPsKeyNum < EACH_PS_KEY_NUM) {
        return threadNumEachPs;
    }

    threadNumEachPs = shardExecMaxCoreNum_ / psSize_;
    if (shardExecMaxCoreNum_ % psSize_ != 0) {
        threadNumEachPs++;
    }

    return threadNumEachPs;
}

HcclResult HdcsOpRemoteLookUpExecutor::FetchUniquedCounts()
{
    for (u32 i = 0; i < psSize_; i++) {
        auto &mappers = conUniqueMappersMgr_.GetMappers();
        EXECEPTION_CATCH(psOffsetAndNum_.Get()[i].count = static_cast<s64>(mappers.at(i).GetUniuqedCount()),
            return HCCL_E_NOT_FOUND);
    }

    return HCCL_SUCCESS;
}

HcclResult HdcsOpRemoteLookUpExecutor::PrepareEnvelopeForRoce(PartitionMapTask *&params,
    void *buffer, u64 count, HcclDataType dataType, u32 psId)
{
    HcclEsRdmaInfoForLookup *rdmaEnveInfos = static_cast<HcclEsRdmaInfoForLookup *>(params->rdmaEnveInfosTransferMem);
    HcclEnvelope &envelope = rdmaEnveInfos[psId].envelope;

    TransData sendData(reinterpret_cast<u64>(buffer), reinterpret_cast<u64>(nullptr), count, dataType,
        false, params->tableId);
    envelope.transData = sendData;

    void *recvValueBuffer = static_cast<u8 *>(params->valueTransferMem) + psOffsetAndNum_.Get()[psId].offset *
        params->valueItemSize;

    RdmaBuffer *rdmaBuffer = reinterpret_cast<RdmaBuffer *>(envelope.rsv);
    // 保留字段传输recvValue的地址
    rdmaBuffer->addr = reinterpret_cast<u64>(recvValueBuffer);

    HCCL_DEBUG("psId[%u] is rdma, buffer[%llu], recvValueBuffer[%llu], count[%llu]", psId,
        hash<void *>{}(buffer), hash<void *>{}(recvValueBuffer), count);

    return HCCL_SUCCESS;
}

HcclResult HdcsOpRemoteLookUpExecutor::TaskSendKeys(OpTaskPtr &task)
{
    PartitionMapTask *params = dynamic_cast<PartitionMapTask *>(task.get());
    HcclEsRdmaInfoForLookup *rdmaEnveInfos = static_cast<HcclEsRdmaInfoForLookup *>(params->rdmaEnveInfosTransferMem);

    CHK_PTR_NULL(rdmaEnveInfos);

    PRINT_ARRAY((u64 *)(rdmaEnveInfos), sizeof(HcclEsRdmaInfoForLookup) * psSize_ / sizeof(u64),
        "LookupDevRdmaEnveInfos");

    u32 valueBytes = params->enableKeyCounter ? ES_KEY_AND_COUNTER_MEM_BYTES_SIZE : ES_KEY_MEM_BYTES_SIZE;
    HcclDataType dataType = params->enableKeyCounter ? HCCL_DATA_TYPE_INT128 : HCCL_DATA_TYPE_INT64;
    if (params->intZerocpyFlag == ZERO_COPY_USED) {
        CHK_PTR_NULL(params->tableIdAddr);
        params->tableIdAddr = reinterpret_cast<void *>(*(reinterpret_cast<u64 *>(params->tableIdAddr)));
    }
    CHK_PTR_NULL(params->tableIdAddr);
    params->tableId = *(static_cast<u32 *>(params->tableIdAddr));
    for (auto iter = rankId2PsId_.begin(); iter != rankId2PsId_.end(); ++iter) {
        u32 rankId = iter->first;
        u32 psId = iter->second;
        void *keysBuffer = static_cast<s8 *>(params->keyTransferMem) + psOffsetAndNum_[psId].offset * valueBytes;
        s64 keyCount = psOffsetAndNum_[psId].count;

        CHK_PRT_RET(UNLIKELY(keyCount > INT32_MAX),
            HCCL_ERROR("keyCount[%lld] > INT32_MAX, failed!", keyCount), HCCL_E_INTERNAL);

        HCCL_DEBUG("tag[%d] tableId[%u] psId[%u] keyCount[%llu] haveRdmaConn[%d] commType[%d]",
            params->tag, params->tableId, psId, static_cast<u64>(keyCount), params->haveRdmaConn,
            rdmaEnveInfos[psId].commType);

        if (params->haveRdmaConn && rdmaEnveInfos[psId].commType == HcclHeterogCommType::RDMA) {
            CHK_RET(PrepareEnvelopeForRoce(params, keysBuffer, static_cast<u64>(keyCount), dataType, psId));
            HCCL_DEBUG("rankId[%u], keyCount[%llu], psId[%u] is rdma", rankId, static_cast<u64>(keyCount), psId);

            // 确保errorStatus正常
            rdmaEnveInfos[psId].errorStatus = 0;
        } else if (!params->haveRdmaConn ||
            (params->haveRdmaConn && rdmaEnveInfos[psId].commType == HcclHeterogCommType::PCIE)) {
            CHK_RET(Send(params, rankId, keysBuffer, static_cast<u64>(keyCount), dataType, requestArray_));
            HCCL_DEBUG("rankId[%u], keyCount[%llu], psId[%u] is pcie", rankId, static_cast<u64>(keyCount), psId);
        } else {
            HCCL_ERROR("psId[%u] commType[%d] is invalid", psId, rdmaEnveInfos[psId].commType);
            return HCCL_E_INTERNAL;
        }
    }

    CHK_RET(WaitSendKeyFinish(task));

    HCCL_INFO("CollRemoteLookUp SendKeys enableKeyCounter[%d], tag[%d], tableId[%u]",
        params->enableKeyCounter, params->tag, params->tableId);

    return HCCL_SUCCESS;
}

HcclResult HdcsOpRemoteLookUpExecutor::WaitSendKeyFinish(OpTaskPtr &task)
{
    CHK_RET(WaitSomeAll(requestArray_));
    requestArray_.clear();
    return HCCL_SUCCESS;
}

HcclResult HdcsOpRemoteLookUpExecutor::Send(PartitionMapTask *params, u32 rankId,
    void *buffer, u64 count, HcclDataType dataType, std::vector<HcclRequest> &requestArray) const
{
    int ret = 0;
    HcclRequest request = nullptr;

    CHK_PTR_NULL(params->comm);
    HcclCommRpc* hcclComm = static_cast<hccl::HcclCommRpc *>(params->comm);
    HcclUserRequire userRequire(params->tableId);
    ret = hcclComm->Isend(buffer, count, dataType, rankId, params->tag, request, userRequire);
    if (ret != HCCL_SUCCESS) {
        HCCL_ERROR("[HdcsOpRemoteLookUpExecutor][Send]HcclIsend failed!!!");
        return HCCL_E_AGAIN;
    }
    requestArray.push_back(request);
    return HCCL_SUCCESS;
}

HcclResult HdcsOpRemoteLookUpExecutor::WaitSomeAll(std::vector<HcclRequest> &requestArray)
{
    s32 allComplete = 0;
    s32 count = requestArray.size();
    HcclRequest requests[count];
    std::copy(requestArray.begin(), requestArray.end(), requests);
    while (allComplete != count) {
        int compCount = 0;
        s32 compIndices[count] = {0};
        HcclStatus compStatus[count] = {0};
        CHK_RET(HcclWaitSomeWithLog(count, requests, &compCount, compIndices, compStatus, false));

        allComplete += compCount;
        if (allComplete != count) {
            SaluSleep(LOOP_SLEEP_TIME_US);
        }
    }

    return HCCL_SUCCESS;
}

HcclResult HdcsOpRemoteLookUpExecutor::ParralAicpuRecoverValue(OpTaskPtr &task)
{
    PartitionMapTask *params = dynamic_cast<PartitionMapTask *>(task.get());

    if (params->disableUnique && (params->outZerocpyFlag == ZERO_COPY_USED)) {
        params->value = reinterpret_cast<float *>(*(static_cast<u64 *>(params->value)));
    }

    HCCL_DEBUG("CollRemoteLookUp ParralAicpuRecoverValue disableUnique[%d], outZerocpyFlag[%d]",
        params->disableUnique, params->outZerocpyFlag);

#ifdef CCL_KERNEL
    ParallelFor(shardExecMaxCoreNum_, 1,
        [&task, this](int64_t start, int64_t end) -> HcclResult {
            for (int64_t i = start; i < end; i++) {
                // start当前标识线程索引，不会截断
                u32 tidx = static_cast<u32>(i);
                auto filtered = [tidx, this] (const u32& index) -> bool {
                    u32 val = index % shardExecMaxCoreNum_;

                    return (tidx != val);
                };

                CHK_RET(AicpuRecoverValue(task, filtered));
            }

            return HCCL_SUCCESS;
        });
#endif

    static_cast<void>(params); // 避免宏函数为空时出现unused variable
    PRINT_ARRAY(static_cast<float *>(params->valueTransferMem), params->keyMaxNum * params->valueItemSize /
        sizeof(float), "LookupRecvValueBeforeCPUGather");
    PRINT_ARRAY(static_cast<float *>(params->value), params->keyMaxNum * params->valueItemSize /
        sizeof(float), "LookupValueAfterCPUGather");

    return HCCL_SUCCESS;
}

HcclResult HdcsOpRemoteLookUpExecutor::AicpuRecoverValue(OpTaskPtr &task, function<bool(const u32&)> filtered)
{
    PartitionMapTask *params = dynamic_cast<PartitionMapTask *>(task.get());
    u32 valueDim = params->valueItemSize / sizeof(float);

    float *dst = static_cast<float *>(params->value);
    float *src = static_cast<float *>(params->valueTransferMem);

    for (u32 i = 0; i < params->keyMaxNum; i++) {
        if (filtered(i)) {
            continue;
        }

        s64 idx = static_cast<s64>(mappingMatrix_[i]);
        if (idx >= static_cast<s64>(params->keyMaxNum) || idx < 0) {
            HCCL_ERROR("idx[%lld] of indices idx[%u], < 0 or >= keyMaxNum[%u]", idx, i,
                params->keyMaxNum);
            return HCCL_E_PARA;
        }

        s64 srcIdx = idx * valueDim;
        s32 sRet = FloatHighSpeedMove(dst + i * valueDim, src + srcIdx, valueDim);
        CHK_PRT_RET(sRet != EOK, HCCL_ERROR("[FloatHighSpeedMove]memcpy failed valueDim[%u]", valueDim),
            HCCL_E_SYSCALL);
    }

    return HCCL_SUCCESS;
}

HcclResult HdcsOpRemoteLookUpExecutor::TaskRecoverValueShard(OpTaskPtr &task)
{
    return HCCL_SUCCESS;
}

HcclResult HdcsOpRemoteLookUpExecutor::TaskRestUnique(OpTaskPtr &task)
{
    PartitionMapTask *params = dynamic_cast<PartitionMapTask *>(task.get());

    if (params->disableUnique) {
        HCCL_DEBUG("CollRemoteLookUp TaskRestUnique disableUnique[%d]", params->disableUnique);
        return HCCL_SUCCESS;
    }

    if (params->cubeIndex <= 0) {
        HCCL_DEBUG("CollRemoteLookUp TaskRestUnique, params->uniqued[%d]", params->uniqued);
        // 清除去重mappingInfo
        CHK_RET(conUniqueMappersMgr_.ClearData());
    }

    return HCCL_SUCCESS;
}

void HdcsOpRemoteLookUpExecutor::PrintLookupRecvSummary(PartitionMapTask *params, u32 psId)
{
    u32 keyBytes = params->enableKeyCounter ? ES_KEY_AND_COUNTER_MEM_BYTES_SIZE : ES_KEY_MEM_BYTES_SIZE;
    void *keysBuffer = static_cast<u8 *>(params->keyTransferMem) +
        psOffsetAndNum_.Get()[psId].offset * keyBytes;
    void *valueBuffer = static_cast<u8 *>(params->valueTransferMem) +
        psOffsetAndNum_.Get()[psId].offset * params->valueItemSize;
    s64 keyCount = psOffsetAndNum_[psId].count;
    u32 valueDim = params->valueItemSize / sizeof(float);
    string printInfo = "lookup recv psId " + std::to_string(psId);
    EsUtils::PrintEsKvDataSummary(printInfo, static_cast<float *>(valueBuffer), keyCount, valueDim,
        static_cast<s64 *>(keysBuffer));
}

HcclResult HdcsOpRemoteLookUpExecutor::TaskWaitLookupFinish(OpTaskPtr &task)
{
    PartitionMapTask *params = dynamic_cast<PartitionMapTask *>(task.get());
    HcclEsRdmaInfoForLookup *rdmaEnveInfos = static_cast<HcclEsRdmaInfoForLookup *>(params->rdmaEnveInfosTransferMem);

    CHK_PTR_NULL(rdmaEnveInfos);

    u32 recvCompleteNum = 0;
    // size == ps count
    unordered_map<u32, s32> recvCompCounts;
    for (auto iter = rankId2PsId_.begin(); iter != rankId2PsId_.end(); ++iter) {
        u32 psId = iter->second;
        if (psOffsetAndNum_.Get()[psId].count == 0) {
            continue;
        }
        recvCompleteNum++;
        recvCompCounts[psId] = 0;
    }

    vector<HcclRequest> requestArray;
    HCCL_INFO("[HdcsOpRemoteLookUpExecutor][WaitLookupFinish] ps size[%d], haveRdmaConn[%d]",
        rankId2PsId_.size(), params->haveRdmaConn);
    while (recvCompleteNum > 0) {
        for (auto iter = rankId2PsId_.begin(); iter != rankId2PsId_.end(); ++iter) {
            u32 rankId = iter->first;
            u32 psId = iter->second;

            if (psOffsetAndNum_.Get()[psId].count == 0 || recvCompCounts[psId] != 0) {
                continue;
            }

            if (params->haveRdmaConn && rdmaEnveInfos[psId].commType == HcclHeterogCommType::RDMA) {
                HCCL_DEBUG("psId[%u] comm is Rdma", psId);
                recvCompleteNum--;
                continue;
            }

            pair<u32, u32> rankIdPsIdPair = make_pair(rankId, psId);
            CHK_RET(RecvRespones(params, rankIdPsIdPair, requestArray, recvCompCounts));
            if (recvCompCounts[psId] != 0) {
                recvCompleteNum--;
            }
        }

        if (recvCompleteNum > 0) {
            SaluSleep(LOOP_IMPROBE_SLEEP_TIME_US);
        }
    }

    CHK_RET(WaitSomeAll(requestArray));

    for (auto iter = rankId2PsId_.begin(); iter != rankId2PsId_.end(); ++iter) {
        u32 rankId = iter->first;
        u32 psId = iter->second;

#ifdef ES_DATA_DFX
    PrintLookupRecvSummary(params, psId);
#endif

        if (params->haveRdmaConn && rdmaEnveInfos[psId].commType == HcclHeterogCommType::RDMA) {
            recvCompCounts[psId] = sizeof(SERVICE_CANCEL_SIGNAL);
            HCCL_DEBUG("psId[%u] comm is Rdma", psId);
        }

        u8 *signal = static_cast<u8 *>(params->valueTransferMem) +
            psOffsetAndNum_.Get()[psId].offset * params->valueItemSize;
        if (IsCancelSingal(signal, recvCompCounts[psId])) {
            HCCL_ERROR("hdds lookup service cancel singal was received, remote rank id[%u] ps id[%u]", rankId, psId);
            return HCCL_E_INTERNAL;
        }
    }

    PRINT_ARRAY(static_cast<float *>(params->valueTransferMem), params->keyMaxNum * params->valueItemSize /
        sizeof(float), "LookupRecvValueTransferMem");

    return HCCL_SUCCESS;
}

HcclResult HdcsOpRemoteLookUpExecutor::RecvRespones(const PartitionMapTask *params,
    std::pair<u32, u32> &rankIdPsIdPair, std::vector<HcclRequest> &requestArray,
    std::unordered_map<u32, s32> &recvCounts)
{
    int ret = 0;
    u32 rank = rankIdPsIdPair.first;
    u32 psId = rankIdPsIdPair.second;
    int flag = 0;
    HcclMessage msg = nullptr;
    HcclStatus status;
    ret = HcclRPCImprobe(rank, params->tag, params->comm, &flag, &msg,  &status);
    CHK_PRT_RET(ret != 0, HCCL_ERROR("HcclRPCImprobe failed"), HCCL_E_INTERNAL);

    CHK_PTR_NULL(msg);
    HcclMessageInfo* hcclMsg = static_cast<HcclMessageInfo *>(msg);
    hccl::HcclCommRpc* hcclComm = static_cast<hccl::HcclCommRpc *>(hcclMsg->commHandle);

    if (flag == HCCL_IMPROBE_COMPLETED) {
        HcclRequest request = nullptr;
        void *buffer = static_cast<u8 *>(params->valueTransferMem) +
            psOffsetAndNum_.Get()[psId].offset * params->valueItemSize;
        u64 count = static_cast<u64>(psOffsetAndNum_.Get()[psId].count) * params->valueItemSize;
        ret = hcclComm->Imrecv(buffer, count, HCCL_DATA_TYPE_INT8, msg, request);
        CHK_PRT_RET(ret != 0, HCCL_ERROR("Imrecv failed"), HCCL_E_INTERNAL);
        HcclRequestInfo* hcclReq = static_cast<HcclRequestInfo *>(request);
        hcclReq->commHandle = hcclComm;
        requestArray.push_back(request);
        ret = HcclRPCGetCount(&status, HCCL_DATA_TYPE_INT8, &(recvCounts[psId]));
        CHK_PRT_RET(ret != 0, HCCL_ERROR("HcclRPCGetCount failed"), HCCL_E_INTERNAL);
    } else {
        recvCounts[psId] = 0;
    }
    return HCCL_SUCCESS;
}
}
